from torch import nn


class ChatNet(nn.Module):
    def __init__(self, vocab_size, embedding_dims, nhead=8, num_encoder_layers=2, num_decoder_layers=2, device=None):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dims, device=device)
        self.transformer = nn.Transformer(embedding_dims, nhead, num_encoder_layers, num_decoder_layers, 1024,
                                          batch_first=True, device=device)
        self.ln = nn.Linear(embedding_dims, vocab_size, device=device)

    def forward(self, inputs, outputs):
        attn_mask = nn.Transformer.generate_square_subsequent_mask(outputs.shape[1])
        inputs = self.embedding(inputs)
        outputs = self.embedding(outputs)
        result = self.transformer(inputs, outputs, tgt_mask=attn_mask, tgt_is_causal=True)
        return self.ln(result)
